from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
import torch

model_pth = "/home/ubuntu/.cache/modelscope/hub/models/Qwen/Qwen2.5-VL-7B-Instruct"
 
# 加载 Qwen2.5-VL-7B-Instruct
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    model_pth,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
 
# 加载 tokenizer 和 processor
tokenizer = AutoTokenizer.from_pretrained(model_pth)
processor = AutoProcessor.from_pretrained(model_pth)
 
# 允许梯度更新
model.enable_input_require_grads()


from datasets import Dataset
import json

# # 读取数据集
# data_path = "path/to/your/dataset.json"  # 请替换为实际路径
# with open(data_path, 'r') as f:
#     data = json.load(f)
#     train_data = data[:-4]  # 划分数据集，保留最后4个样本作为测试集
#     test_data = data[-4:]

# # 保存数据
# with open("train_data.json", "w") as f:
#     json.dump(train_data, f)
# with open("test_data.json", "w") as f:
#     json.dump(test_data, f)

# 加载数据集
train_ds = Dataset.from_json("data1.json")


from qwen_vl_utils import process_vision_info
import torch

def process_func(example):
    # print(example)
    """
    预处理输入数据
    """
    MAX_LENGTH = 8192
    input_content = example['0']
    output_content = example['1']

    # 构造多模态对话
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": input_content},
            ],
        }
    ]
    
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = processor(
        text=[text],
        padding=True,
        return_tensors="pt",
    )

    inputs = {key: value.tolist() for key, value in inputs.items()}
    
    # 构造目标输出
    response = tokenizer(f"{output_content}", add_special_tokens=False)
    input_ids = inputs["input_ids"][0] + response["input_ids"] + [tokenizer.pad_token_id]
    attention_mask = inputs["attention_mask"][0] + response["attention_mask"] + [1]
    labels = [-100] * len(inputs["input_ids"][0]) + response["input_ids"] + [tokenizer.pad_token_id]

    # 截断
    if len(input_ids) > MAX_LENGTH:
        input_ids = input_ids[:MAX_LENGTH]
        attention_mask = attention_mask[:MAX_LENGTH]
        labels = labels[:MAX_LENGTH]

    return {
        "input_ids": torch.tensor(input_ids),
        "attention_mask": torch.tensor(attention_mask),
        "labels": torch.tensor(labels)
    }


# 处理数据
train_dataset = train_ds.map(process_func)

# 确保数据加载成功
print(f"Train dataset size: {len(train_dataset)}")
print(train_dataset[0])  # 检查数据格式


from peft import LoraConfig, get_peft_model
 
config = LoraConfig(
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    inference_mode=False,
    r=64,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
)
 
# 将 LoRA 应用于模型
peft_model = get_peft_model(model, config)


from transformers import TrainingArguments, Trainer, DataCollatorForSeq2Seq
import os
 
args = TrainingArguments(
    output_dir="output/Qwen2.5-VL-LoRA",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    logging_steps=10,
    num_train_epochs=3,
    save_steps=74,
    learning_rate=1e-4,
    gradient_checkpointing=True,
    report_to="none"
)
 
trainer = Trainer(
    model=peft_model,
    args=args,
    train_dataset=train_dataset,  # 需要提供数据
    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),
)
 
trainer.train()
